wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-10 00:17:16 +08:00
Code Issues Actions 37 Packages Projects Releases Wiki Activity

Make a context manager for cast_bias_weight and use it.

Browse Source
This commit is contained in:
blepping 2026-07-03 08:33:45 -06:00
parent 77917ed3a6
commit 528c44de2d
4 changed files with 118 additions and 158 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
comfy/background_removal/birefnet.py
View File

@ -433,19 +433,16 @@ class DeformableConv2d(nn.Module):
def forward(self, x): def forward(self, x):
offset = self.offset_conv(x) offset = self.offset_conv(x)
modulator = 2. * torch.sigmoid(self.modulator_conv(x)) modulator = 2. * torch.sigmoid(self.modulator_conv(x))
weight, bias, offload_info = comfy.ops.cast_bias_weight(self.regular_conv, x, offloadable=True) with comfy.ops.CastBiasWeightContext(self.regular_conv, x, offloadable=True) as (weight, _bias):
return deform_conv2d(
x = deform_conv2d( input=x,
input=x, offset=offset,
offset=offset, weight=weight,
weight=weight, bias=None,
bias=None, padding=self.padding,
padding=self.padding, mask=modulator,
mask=modulator, stride=self.stride,
stride=self.stride, )
)
comfy.ops.uncast_bias_weight(self.regular_conv, weight, bias, offload_info)
return x
class BasicDecBlk(nn.Module): class BasicDecBlk(nn.Module):
def __init__(self, in_channels=64, out_channels=64, inter_channels=64, device=None, dtype=None, operations=None): def __init__(self, in_channels=64, out_channels=64, inter_channels=64, device=None, dtype=None, operations=None):

11
comfy/controlnet.py
View File

@ -381,13 +381,10 @@ class ControlLoraOps:
self.bias = None self.bias = None
def forward(self, input): def forward(self, input):
weight, bias, offload_stream = comfy.ops.cast_bias_weight(self, input, offloadable=True) with comfy.ops.CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
if self.up is not None: if self.up is None:
x = torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias) return torch.nn.functional.linear(input, weight, bias)
else: return torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
x = torch.nn.functional.linear(input, weight, bias)
comfy.ops.uncast_bias_weight(self, weight, bias, offload_stream)
return x
class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp): class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp):
def __init__( def __init__(

228
comfy/ops.py
View File

@ -402,6 +402,26 @@ def uncast_bias_weight(s, weight, bias, offload_stream):
device = bias_a.device device = bias_a.device
os.wait_stream(comfy.model_management.current_stream(device)) os.wait_stream(comfy.model_management.current_stream(device))
class CastBiasWeightContext:
# 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 not self.slf:
return
slf, state = self.slf, self.state
self.state = self.slf = None
uncast_bias_weight(slf, *state)
class CastWeightBiasOp: class CastWeightBiasOp:
comfy_cast_weights = False comfy_cast_weights = False
@ -490,10 +510,8 @@ class disable_weight_init:
return None return None
def forward_comfy_cast_weights(self, input): def forward_comfy_cast_weights(self, input):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = torch.nn.functional.linear(input, weight, bias) return torch.nn.functional.linear(input, weight, bias)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -507,10 +525,8 @@ class disable_weight_init:
return None return None
def forward_comfy_cast_weights(self, input): def forward_comfy_cast_weights(self, input):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = self._conv_forward(input, weight, bias) return self._conv_forward(input, weight, bias)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -524,10 +540,8 @@ class disable_weight_init:
return None return None
def forward_comfy_cast_weights(self, input): def forward_comfy_cast_weights(self, input):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = self._conv_forward(input, weight, bias) return self._conv_forward(input, weight, bias)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -552,10 +566,8 @@ class disable_weight_init:
return super()._conv_forward(input, weight, bias, *args, **kwargs) return super()._conv_forward(input, weight, bias, *args, **kwargs)
def forward_comfy_cast_weights(self, input, autopad=None): def forward_comfy_cast_weights(self, input, autopad=None):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = self._conv_forward(input, weight, bias, autopad=autopad) return self._conv_forward(input, weight, bias, autopad=autopad)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -569,10 +581,8 @@ class disable_weight_init:
return None return None
def forward_comfy_cast_weights(self, input): def forward_comfy_cast_weights(self, input):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps) return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -586,12 +596,10 @@ class disable_weight_init:
return None return None
def forward_comfy_cast_weights(self, input): def forward_comfy_cast_weights(self, input):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) 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_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 running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None
x = torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps) return torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -605,15 +613,8 @@ class disable_weight_init:
return None return None
def forward_comfy_cast_weights(self, input): def forward_comfy_cast_weights(self, input):
if self.weight is not None: with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
else:
weight = None
bias = None
offload_stream = None
x = torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -628,15 +629,8 @@ class disable_weight_init:
return None return None
def forward_comfy_cast_weights(self, input): def forward_comfy_cast_weights(self, input):
if self.weight is not None: with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) return torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
else:
weight = None
bias = None
offload_stream = None
x = torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -655,12 +649,10 @@ class disable_weight_init:
input, output_size, self.stride, self.padding, self.kernel_size, input, output_size, self.stride, self.padding, self.kernel_size,
num_spatial_dims, self.dilation) num_spatial_dims, self.dilation)
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = torch.nn.functional.conv_transpose2d( return torch.nn.functional.conv_transpose2d(
input, weight, bias, self.stride, self.padding, input, weight, bias, self.stride, self.padding,
output_padding, self.groups, self.dilation) output_padding, self.groups, self.dilation)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -679,12 +671,10 @@ class disable_weight_init:
input, output_size, self.stride, self.padding, self.kernel_size, input, output_size, self.stride, self.padding, self.kernel_size,
num_spatial_dims, self.dilation) num_spatial_dims, self.dilation)
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = torch.nn.functional.conv_transpose1d( return torch.nn.functional.conv_transpose1d(
input, weight, bias, self.stride, self.padding, input, weight, bias, self.stride, self.padding,
output_padding, self.groups, self.dilation) output_padding, self.groups, self.dilation)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -749,10 +739,8 @@ class disable_weight_init:
output_dtype = out_dtype output_dtype = out_dtype
if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16: if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
out_dtype = None out_dtype = None
weight, bias, offload_stream = cast_bias_weight(self, device=input.device, dtype=out_dtype, offloadable=True) with CastBiasWeightContext(self, device=input.device, dtype=out_dtype, offloadable=True) as (weight, bias):
x = 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) 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)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
@ -828,7 +816,6 @@ def fp8_linear(self, input):
if input.ndim != 2: if input.ndim != 2:
return None return None
lora_compute_dtype=comfy.model_management.lora_compute_dtype(input.device) lora_compute_dtype=comfy.model_management.lora_compute_dtype(input.device)
w, bias, offload_stream = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True)
scale_weight = torch.ones((), device=input.device, dtype=torch.float32) scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
scale_input = torch.ones((), device=input.device, dtype=torch.float32) scale_input = torch.ones((), device=input.device, dtype=torch.float32)
@ -837,15 +824,16 @@ def fp8_linear(self, input):
layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape)) layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input) quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
# Wrap weight in QuantizedTensor - this enables unified dispatch with CastBiasWeightContext(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True) as (w, bias):
# Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py! # Wrap weight in QuantizedTensor - this enables unified dispatch
layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape)) # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight) w_shape = tuple(w.shape)
o = torch.nn.functional.linear(quantized_input, quantized_weight, bias) layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=w_shape)
quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
uncast_bias_weight(self, w, bias, offload_stream)
if tensor_3d: if tensor_3d:
o = o.reshape((input_shape[0], input_shape[1], w.shape[0])) o = o.reshape((input_shape[0], input_shape[1], w_shape[0]))
return o return o
@ -865,10 +853,8 @@ class fp8_ops(manual_cast):
except Exception as e: except Exception as e:
logging.info("Exception during fp8 op: {}".format(e)) logging.info("Exception during fp8 op: {}".format(e))
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = torch.nn.functional.linear(input, weight, bias) return torch.nn.functional.linear(input, weight, bias)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
CUBLAS_IS_AVAILABLE = False CUBLAS_IS_AVAILABLE = False
try: try:
@ -884,10 +870,8 @@ if CUBLAS_IS_AVAILABLE:
return None return None
def forward_comfy_cast_weights(self, input): def forward_comfy_cast_weights(self, input):
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
x = cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias) return cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
uncast_bias_weight(self, weight, bias, offload_stream)
return x
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
run_every_op() run_every_op()
@ -1207,29 +1191,28 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
want_requant=False, want_requant=False,
weight_only_quant=False, weight_only_quant=False,
): ):
if weight_only_quant: if not weight_only_quant:
weight, bias, offload_stream = cast_bias_weight( with CastBiasWeightContext(
self,
input=None,
dtype=self.weight.dtype,
device=input.device,
bias_dtype=input.dtype,
offloadable=True,
compute_dtype=compute_dtype,
want_requant=True,
)
weight = weight.to(dtype=input.dtype)
else:
weight, bias, offload_stream = cast_bias_weight(
self, self,
input, input,
offloadable=True, offloadable=True,
compute_dtype=compute_dtype, compute_dtype=compute_dtype,
want_requant=want_requant, want_requant=want_requant,
) ) as (weight, bias):
x = self._forward(input, weight, bias) return self._forward(input, weight, bias)
uncast_bias_weight(self, weight, bias, offload_stream)
return x with CastBiasWeightContext(
self,
input=None,
dtype=self.weight.dtype,
device=input.device,
bias_dtype=input.dtype,
offloadable=True,
compute_dtype=compute_dtype,
want_requant=True,
) as (weight, bias):
weight = weight.to(dtype=input.dtype)
return self._forward(input, weight, bias)
def forward(self, input, *args, **kwargs): def forward(self, input, *args, **kwargs):
run_every_op() run_every_op()
@ -1249,25 +1232,20 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
# Training path: quantized forward with compute_dtype backward via autograd function # Training path: quantized forward with compute_dtype backward via autograd function
if (input.requires_grad and _use_quantized and quantize_input): if (input.requires_grad and _use_quantized and quantize_input):
with CastBiasWeightContext(
weight, bias, offload_stream = cast_bias_weight(
self, self,
input, input,
offloadable=True, offloadable=True,
compute_dtype=compute_dtype, compute_dtype=compute_dtype,
want_requant=True want_requant=True
) ) as (weight, bias):
scale = getattr(self, 'input_scale', None)
if scale is not None:
scale = comfy.model_management.cast_to_device(scale, input.device, None)
scale = getattr(self, 'input_scale', None) return QuantLinearFunc.apply(
if scale is not None: input, weight, bias, self.layout_type, scale, compute_dtype
scale = comfy.model_management.cast_to_device(scale, input.device, None) )
output = QuantLinearFunc.apply(
input, weight, bias, self.layout_type, scale, compute_dtype
)
uncast_bias_weight(self, weight, bias, offload_stream)
return output
# Inference path (unchanged) # Inference path (unchanged)
if _use_quantized and quantize_input: if _use_quantized and quantize_input:
@ -1378,13 +1356,11 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
"""Cast the whole bank once; expert_linear inside reuses the cast. """Cast the whole bank once; expert_linear inside reuses the cast.
Not re-entrant — do not nest calls on the same instance. Not re-entrant — do not nest calls on the same instance.
""" """
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as self._resident_bank:
self._resident_bank = (weight, bias) try:
try: yield self
yield self finally:
finally: self._resident_bank = None
self._resident_bank = None
uncast_bias_weight(self, weight, bias, offload_stream)
def expert_linear(self, input: torch.Tensor, i: int) -> torch.Tensor: def expert_linear(self, input: torch.Tensor, i: int) -> torch.Tensor:
"""Linear against expert i's weight (with optional bias).""" """Linear against expert i's weight (with optional bias)."""
@ -1392,11 +1368,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
if resident is not None: if resident is not None:
weight, bias = resident weight, bias = resident
return self._expert_linear_impl(input, weight, bias, i) return self._expert_linear_impl(input, weight, bias, i)
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True) with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
try:
return self._expert_linear_impl(input, weight, bias, i) return self._expert_linear_impl(input, weight, bias, i)
finally:
uncast_bias_weight(self, weight, bias, offload_stream)
def _expert_linear_impl(self, input, weight, bias, i): def _expert_linear_impl(self, input, weight, bias, i):
if isinstance(weight, QuantizedTensor): if isinstance(weight, QuantizedTensor):
@ -1487,17 +1460,16 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
# Optimized path: lookup in fp8, dequantize only the selected rows. # Optimized path: lookup in fp8, dequantize only the selected rows.
if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0: if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0:
qdata, _, offload_stream = cast_bias_weight(self, device=input.device, dtype=weight.dtype, offloadable=True) with CastBiasWeightContext(self, device=input.device, dtype=weight.dtype, offloadable=True) as (qdata, _bias):
if isinstance(qdata, QuantizedTensor): if isinstance(qdata, QuantizedTensor):
scale = qdata._params.scale scale = qdata._params.scale
qdata = qdata._qdata qdata = qdata._qdata
else: else:
scale = None scale = None
x = torch.nn.functional.embedding( x = torch.nn.functional.embedding(
input, qdata, self.padding_idx, self.max_norm, input, qdata, self.padding_idx, self.max_norm,
self.norm_type, self.scale_grad_by_freq, self.sparse) self.norm_type, self.scale_grad_by_freq, self.sparse)
uncast_bias_weight(self, qdata, None, offload_stream)
target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
x = x.to(dtype=target_dtype) x = x.to(dtype=target_dtype)
if scale is not None and scale != 1.0: if scale is not None and scale != 1.0:

14
comfy/text_encoders/llama.py
View File

@ -859,16 +859,10 @@ class BaseGenerate:
else: else:
module = self.model.embed_tokens module = self.model.embed_tokens
offload_stream = None if not module.comfy_cast_weights:
if module.comfy_cast_weights: return torch.nn.functional.linear(input, self.model.embed_tokens.weight.to(x), None)
weight, _, offload_stream = comfy.ops.cast_bias_weight(module, input, offloadable=True) with comfy.ops.CastBiasWeightContext(module, input, offloadable=True) as (weight, _bias):
else: return torch.nn.functional.linear(input, weight, None)
weight = self.model.embed_tokens.weight.to(x)
x = torch.nn.functional.linear(input, weight, None)
comfy.ops.uncast_bias_weight(module, weight, None, offload_stream)
return x
def init_kv_cache(self, batch, max_cache_len, device, execution_dtype): def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
model_config = self.model.config model_config = self.model.config