Merge branch 'master' into fix/empty-progress-log
This commit is contained in:
Alexis Rolland
GitHub
commit
9692e435b2
|
|
@ -433,19 +433,16 @@ class DeformableConv2d(nn.Module):
|
|||
def forward(self, x):
|
||||
offset = self.offset_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)
|
||||
|
||||
x = deform_conv2d(
|
||||
input=x,
|
||||
offset=offset,
|
||||
weight=weight,
|
||||
bias=None,
|
||||
padding=self.padding,
|
||||
mask=modulator,
|
||||
stride=self.stride,
|
||||
)
|
||||
comfy.ops.uncast_bias_weight(self.regular_conv, weight, bias, offload_info)
|
||||
return x
|
||||
with comfy.ops.CastBiasWeightContext(self.regular_conv, x, offloadable=True) as (weight, _bias):
|
||||
return deform_conv2d(
|
||||
input=x,
|
||||
offset=offset,
|
||||
weight=weight,
|
||||
bias=None,
|
||||
padding=self.padding,
|
||||
mask=modulator,
|
||||
stride=self.stride,
|
||||
)
|
||||
|
||||
class BasicDecBlk(nn.Module):
|
||||
def __init__(self, in_channels=64, out_channels=64, inter_channels=64, device=None, dtype=None, operations=None):
|
||||
|
|
|
|||
|
|
@ -381,13 +381,10 @@ class ControlLoraOps:
|
|||
self.bias = None
|
||||
|
||||
def forward(self, input):
|
||||
weight, bias, offload_stream = comfy.ops.cast_bias_weight(self, input, offloadable=True)
|
||||
if self.up is not 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)
|
||||
else:
|
||||
x = torch.nn.functional.linear(input, weight, bias)
|
||||
comfy.ops.uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with comfy.ops.CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
if self.up is None:
|
||||
return torch.nn.functional.linear(input, weight, bias)
|
||||
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)
|
||||
|
||||
class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp):
|
||||
def __init__(
|
||||
|
|
|
|||
241
comfy/ops.py
241
comfy/ops.py
|
|
@ -452,6 +452,26 @@ def uncast_bias_weight(s, weight, bias, offload_stream):
|
|||
device = bias_a.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 self.slf is None:
|
||||
return
|
||||
slf, state = self.slf, self.state
|
||||
self.state = self.slf = None
|
||||
uncast_bias_weight(slf, *state)
|
||||
|
||||
class CastWeightBiasOp:
|
||||
comfy_cast_weights = False
|
||||
|
|
@ -538,10 +558,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.linear(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
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()
|
||||
|
|
@ -555,10 +573,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = self._conv_forward(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return self._conv_forward(input, weight, bias)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -572,10 +588,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = self._conv_forward(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return self._conv_forward(input, weight, bias)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -600,10 +614,8 @@ class disable_weight_init:
|
|||
return super()._conv_forward(input, weight, bias, *args, **kwargs)
|
||||
|
||||
def forward_comfy_cast_weights(self, input, autopad=None):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = self._conv_forward(input, weight, bias, autopad=autopad)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
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()
|
||||
|
|
@ -617,10 +629,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
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()
|
||||
|
|
@ -634,12 +644,10 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
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
|
||||
x = 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
|
||||
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()
|
||||
|
|
@ -653,15 +661,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
if self.weight is not None:
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
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
|
||||
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()
|
||||
|
|
@ -676,15 +677,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
if self.weight is not None:
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
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
|
||||
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()
|
||||
|
|
@ -703,12 +697,10 @@ class disable_weight_init:
|
|||
input, output_size, self.stride, self.padding, self.kernel_size,
|
||||
num_spatial_dims, self.dilation)
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.conv_transpose2d(
|
||||
input, weight, bias, self.stride, self.padding,
|
||||
output_padding, self.groups, self.dilation)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
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()
|
||||
|
|
@ -727,12 +719,10 @@ class disable_weight_init:
|
|||
input, output_size, self.stride, self.padding, self.kernel_size,
|
||||
num_spatial_dims, self.dilation)
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.conv_transpose1d(
|
||||
input, weight, bias, self.stride, self.padding,
|
||||
output_padding, self.groups, self.dilation)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
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()
|
||||
|
|
@ -795,10 +785,8 @@ class disable_weight_init:
|
|||
output_dtype = out_dtype
|
||||
if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
|
||||
out_dtype = None
|
||||
weight, bias, offload_stream = cast_bias_weight(self, device=input.device, dtype=out_dtype, offloadable=True)
|
||||
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)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
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):
|
||||
|
|
@ -874,7 +862,6 @@ def fp8_linear(self, input):
|
|||
if input.ndim != 2:
|
||||
return None
|
||||
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_input = torch.ones((), device=input.device, dtype=torch.float32)
|
||||
|
|
@ -883,15 +870,16 @@ def fp8_linear(self, input):
|
|||
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)
|
||||
|
||||
# Wrap weight in QuantizedTensor - this enables unified dispatch
|
||||
# Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
|
||||
layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
|
||||
quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
|
||||
o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
|
||||
with CastBiasWeightContext(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True) as (w, bias):
|
||||
# Wrap weight in QuantizedTensor - this enables unified dispatch
|
||||
# Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
|
||||
w_shape = tuple(w.shape)
|
||||
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:
|
||||
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
|
||||
|
||||
|
|
@ -911,10 +899,8 @@ class fp8_ops(manual_cast):
|
|||
except Exception as e:
|
||||
logging.info("Exception during fp8 op: {}".format(e))
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.linear(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.linear(input, weight, bias)
|
||||
|
||||
CUBLAS_IS_AVAILABLE = False
|
||||
try:
|
||||
|
|
@ -930,10 +916,8 @@ if CUBLAS_IS_AVAILABLE:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -1344,29 +1328,28 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
want_requant=False,
|
||||
weight_only_quant=False,
|
||||
):
|
||||
if weight_only_quant:
|
||||
weight, bias, offload_stream = cast_bias_weight(
|
||||
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(
|
||||
if not weight_only_quant:
|
||||
with CastBiasWeightContext(
|
||||
self,
|
||||
input,
|
||||
offloadable=True,
|
||||
compute_dtype=compute_dtype,
|
||||
want_requant=want_requant,
|
||||
)
|
||||
x = self._forward(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
) as (weight, bias):
|
||||
return self._forward(input, weight, bias)
|
||||
|
||||
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):
|
||||
run_every_op()
|
||||
|
|
@ -1391,25 +1374,20 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
|
||||
# Training path: quantized forward with compute_dtype backward via autograd function
|
||||
if (input.requires_grad and _use_quantized and quantize_input):
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(
|
||||
with CastBiasWeightContext(
|
||||
self,
|
||||
input,
|
||||
offloadable=True,
|
||||
compute_dtype=compute_dtype,
|
||||
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)
|
||||
if scale is not None:
|
||||
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
|
||||
return QuantLinearFunc.apply(
|
||||
input, weight, bias, self.layout_type, scale, compute_dtype
|
||||
)
|
||||
|
||||
# Inference path (unchanged)
|
||||
if _use_quantized and quantize_input:
|
||||
|
|
@ -1520,13 +1498,11 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
"""Cast the whole bank once; expert_linear inside reuses the cast.
|
||||
Not re-entrant — do not nest calls on the same instance.
|
||||
"""
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
self._resident_bank = (weight, bias)
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._resident_bank = None
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as self._resident_bank:
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._resident_bank = None
|
||||
|
||||
def expert_linear(self, input: torch.Tensor, i: int) -> torch.Tensor:
|
||||
"""Linear against expert i's weight (with optional bias)."""
|
||||
|
|
@ -1534,11 +1510,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
if resident is not None:
|
||||
weight, bias = resident
|
||||
return self._expert_linear_impl(input, weight, bias, i)
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
try:
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
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):
|
||||
if isinstance(weight, QuantizedTensor):
|
||||
|
|
@ -1641,25 +1614,23 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
|
||||
# Optimized path: lookup in fp8/int8, dequantize only the selected rows.
|
||||
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)
|
||||
if isinstance(qdata, QuantizedTensor):
|
||||
params = qdata._params
|
||||
scale = params.scale
|
||||
qdata = qdata._qdata
|
||||
else:
|
||||
params = weight._params
|
||||
scale = None
|
||||
with CastBiasWeightContext(self, device=input.device, dtype=weight.dtype, offloadable=True) as (qdata, _bias):
|
||||
if isinstance(qdata, QuantizedTensor):
|
||||
params = qdata._params
|
||||
scale = params.scale
|
||||
qdata = qdata._qdata
|
||||
else:
|
||||
params = weight._params
|
||||
scale = None
|
||||
|
||||
# int8: per-row scale possible ConvRot, so let the layout do the gather
|
||||
if self.quant_format == "int8_tensorwise":
|
||||
x = get_layout_class(self.layout_type).dequantize_embedding(qdata, params, input)
|
||||
uncast_bias_weight(self, qdata, None, offload_stream)
|
||||
return x if out_dtype is None else x.to(dtype=out_dtype)
|
||||
# int8: per-row scale possible ConvRot, so let the layout do the gather
|
||||
if self.quant_format == "int8_tensorwise":
|
||||
x = get_layout_class(self.layout_type).dequantize_embedding(qdata, params, input)
|
||||
return x if out_dtype is None else x.to(dtype=out_dtype)
|
||||
|
||||
x = torch.nn.functional.embedding(
|
||||
input, qdata, self.padding_idx, self.max_norm,
|
||||
self.norm_type, self.scale_grad_by_freq, self.sparse)
|
||||
uncast_bias_weight(self, qdata, None, offload_stream)
|
||||
x = torch.nn.functional.embedding(
|
||||
input, qdata, self.padding_idx, self.max_norm,
|
||||
self.norm_type, self.scale_grad_by_freq, self.sparse)
|
||||
target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
|
||||
x = x.to(dtype=target_dtype)
|
||||
if scale is not None and scale != 1.0:
|
||||
|
|
|
|||
|
|
@ -868,16 +868,10 @@ class BaseGenerate:
|
|||
else:
|
||||
module = self.model.embed_tokens
|
||||
|
||||
offload_stream = None
|
||||
if module.comfy_cast_weights:
|
||||
weight, _, offload_stream = comfy.ops.cast_bias_weight(module, input, offloadable=True)
|
||||
else:
|
||||
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
|
||||
if not module.comfy_cast_weights:
|
||||
return torch.nn.functional.linear(input, self.model.embed_tokens.weight.to(x), None)
|
||||
with comfy.ops.CastBiasWeightContext(module, input, offloadable=True) as (weight, _bias):
|
||||
return torch.nn.functional.linear(input, weight, None)
|
||||
|
||||
def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
|
||||
model_config = self.model.config
|
||||
|
|
|
|||
|
|
@ -197,6 +197,7 @@ def normalize_queue_item(item: tuple, status: str) -> dict:
|
|||
'priority': priority,
|
||||
'create_time': create_time,
|
||||
'outputs_count': 0,
|
||||
'previewable_outputs_count': 0,
|
||||
'workflow_id': workflow_id,
|
||||
})
|
||||
|
||||
|
|
@ -215,6 +216,7 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
|
|||
|
||||
outputs = history_item.get('outputs', {})
|
||||
outputs_count, preview_output = get_outputs_summary(outputs)
|
||||
previewable_outputs_count = count_previewable_outputs(outputs)
|
||||
|
||||
execution_error = None
|
||||
execution_start_time = None
|
||||
|
|
@ -251,6 +253,7 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
|
|||
'execution_end_time': execution_end_time,
|
||||
'execution_error': execution_error,
|
||||
'outputs_count': outputs_count,
|
||||
'previewable_outputs_count': previewable_outputs_count,
|
||||
'preview_output': preview_output,
|
||||
'workflow_id': workflow_id,
|
||||
})
|
||||
|
|
@ -345,6 +348,33 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
|
|||
return count, preview_output or fallback_preview or text_file_fallback or text_fallback
|
||||
|
||||
|
||||
def count_previewable_outputs(outputs: dict) -> int:
|
||||
"""
|
||||
Count only outputs that would actually render in the expanded asset view,
|
||||
i.e. items is_previewable() accepts (image/video/audio/3D/text). Kept
|
||||
separate from get_outputs_summary()'s outputs_count, which counts every
|
||||
output item regardless of media type, so a job with a non-previewable
|
||||
saved file alongside real media (e.g. SaveLatent's .latent output next to
|
||||
a SaveImage output) doesn't inflate the Media Assets badge beyond what
|
||||
the expanded view shows.
|
||||
"""
|
||||
count = 0
|
||||
for node_outputs in outputs.values():
|
||||
if not isinstance(node_outputs, dict):
|
||||
continue
|
||||
for media_type, items in node_outputs.items():
|
||||
if media_type == 'animated' or not isinstance(items, list):
|
||||
continue
|
||||
for item in items:
|
||||
if not isinstance(item, dict):
|
||||
item = normalize_output_item(item)
|
||||
if item is None:
|
||||
continue
|
||||
if is_previewable(media_type, item):
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def apply_sorting(jobs: list[dict], sort_by: str, sort_order: str) -> list[dict]:
|
||||
"""Sort jobs list by specified field and order."""
|
||||
reverse = (sort_order == 'desc')
|
||||
|
|
|
|||
|
|
@ -486,7 +486,10 @@ class ImageCompositor(io.ComfyNode):
|
|||
node_id="ImageCompositor",
|
||||
display_name="Create Layered Image",
|
||||
category="image",
|
||||
search_aliases=["compositor", "composite", "layer", "layers", "layer editor", "psd"],
|
||||
is_experimental=True,
|
||||
# both flags on purpose: terminal compositor graphs must execute (the
|
||||
# editor needs a run to open), and cache hits must replay the layer UI
|
||||
is_output_node=True,
|
||||
has_intermediate_output=True,
|
||||
inputs=[
|
||||
|
|
@ -605,7 +608,7 @@ class AddLayer(io.ComfyNode):
|
|||
options=list(_LAYER_MODES),
|
||||
default="normal",
|
||||
optional=True,
|
||||
tooltip="Initial blend mode.",
|
||||
tooltip="Initial blend mode, applied against the layers below. On the bottom layer over the default transparent background, non-normal modes produce transparency.",
|
||||
),
|
||||
io.Float.Input(
|
||||
"rotation",
|
||||
|
|
|
|||
|
|
@ -591,7 +591,7 @@ class SamplerER_SDE(io.ComfyNode):
|
|||
inputs=[
|
||||
io.Combo.Input("solver_type", options=["ER-SDE", "Reverse-time SDE", "ODE"]),
|
||||
io.Int.Input("max_stage", default=3, min=1, max=3, advanced=True),
|
||||
io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, tooltip="Stochastic strength of reverse-time SDE.\nWhen eta=0, it reduces to deterministic ODE. This setting doesn't apply to ER-SDE solver type.", advanced=True),
|
||||
io.Float.Input("eta", default=1.0, min=0.0, max=10.0, step=0.01, round=False, tooltip="Stochastic strength of SDEs.\nWhen eta=0, they reduce to deterministic ODE.\nLarge eta may cause invalid outputs. If this occurs, try decreasing this value.", advanced=True),
|
||||
io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
|
||||
],
|
||||
outputs=[io.Sampler.Output()]
|
||||
|
|
@ -599,21 +599,35 @@ class SamplerER_SDE(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, solver_type, max_stage, eta, s_noise) -> io.NodeOutput:
|
||||
if solver_type == "ODE" or (solver_type == "Reverse-time SDE" and eta == 0):
|
||||
eta = 0
|
||||
s_noise = 0
|
||||
# Extend existing noise scalers phi(x) with eta-controlled noise scalers:
|
||||
# psi(x) = x**(1-eta) * phi(x)**eta
|
||||
# where eta is constant and directly scales the h^2(t) contribution.
|
||||
|
||||
def reverse_time_sde_noise_scaler(x):
|
||||
def er_sde_noise_scaler(x: torch.Tensor) -> torch.Tensor:
|
||||
return x * ((x ** 0.3).exp() + 10.0) ** eta
|
||||
|
||||
def reverse_time_sde_noise_scaler(x: torch.Tensor) -> torch.Tensor:
|
||||
return x ** (eta + 1)
|
||||
|
||||
if solver_type == "ER-SDE":
|
||||
# Use the default one in sample_er_sde()
|
||||
noise_scaler = None
|
||||
else:
|
||||
noise_scaler = reverse_time_sde_noise_scaler
|
||||
def ode_noise_scaler(x: torch.Tensor) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
solver_scalers = {
|
||||
"ER-SDE": er_sde_noise_scaler,
|
||||
"Reverse-time SDE": reverse_time_sde_noise_scaler,
|
||||
"ODE": ode_noise_scaler,
|
||||
}
|
||||
|
||||
if solver_type == "ODE" or eta == 0:
|
||||
s_noise = 0.0
|
||||
solver_type = "ODE"
|
||||
noise_scaler = solver_scalers[solver_type]
|
||||
|
||||
sampler_name = "er_sde"
|
||||
sampler = comfy.samplers.ksampler(sampler_name, {"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage})
|
||||
sampler = comfy.samplers.ksampler(
|
||||
sampler_name,
|
||||
{"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage},
|
||||
)
|
||||
return io.NodeOutput(sampler)
|
||||
|
||||
get_sampler = execute
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ class ImageUpscaleWithModel(io.ComfyNode):
|
|||
|
||||
memory_required = (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
|
||||
memory_required += image.nelement() * image.element_size()
|
||||
model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required)
|
||||
model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required, force_full_load=True)
|
||||
|
||||
in_img = image.movedim(-1,-3).to(device)
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from comfy_execution.jobs import (
|
|||
normalize_output_item,
|
||||
normalize_outputs,
|
||||
get_outputs_summary,
|
||||
count_previewable_outputs,
|
||||
apply_sorting,
|
||||
has_3d_extension,
|
||||
validate_job_id,
|
||||
|
|
@ -361,6 +362,79 @@ class TestGetOutputsSummary:
|
|||
assert preview['mediaType'] == 'files'
|
||||
|
||||
|
||||
class TestCountPreviewableOutputs:
|
||||
"""Unit tests for count_previewable_outputs()
|
||||
|
||||
Kept separate from get_outputs_summary()'s outputs_count: the Media Assets
|
||||
badge should reflect only what the expanded asset view actually renders
|
||||
(previewable outputs), while outputs_count keeps counting every output
|
||||
item for other consumers.
|
||||
"""
|
||||
|
||||
def test_empty_outputs(self):
|
||||
assert count_previewable_outputs({}) == 0
|
||||
|
||||
def test_previewable_outputs_all_counted(self):
|
||||
"""When every output is previewable, the two counts should match."""
|
||||
outputs = {
|
||||
'node1': {'images': [{'filename': 'a.png', 'type': 'output'}]},
|
||||
'node2': {'images': [{'filename': 'b.png', 'type': 'output'}]},
|
||||
}
|
||||
outputs_count, _ = get_outputs_summary(outputs)
|
||||
assert count_previewable_outputs(outputs) == outputs_count == 2
|
||||
|
||||
def test_save_latent_counted_but_not_previewable(self):
|
||||
"""SaveLatent (nodes.py) emits a real saved file under the 'latents'
|
||||
media type: {'latents': [{'filename': '..._00001_.latent',
|
||||
'subfolder': '', 'type': 'output'}]}. It has no previewable media
|
||||
type, format, or extension, so it inflates outputs_count without
|
||||
ever rendering in the expanded asset view."""
|
||||
outputs = {
|
||||
'node1': {
|
||||
'images': [{'filename': 'ComfyUI_00001_.png', 'subfolder': '', 'type': 'output'}]
|
||||
},
|
||||
'node2': {
|
||||
'latents': [{'filename': 'ComfyUI_00001_.latent', 'subfolder': '', 'type': 'output'}]
|
||||
},
|
||||
}
|
||||
outputs_count, _ = get_outputs_summary(outputs)
|
||||
assert outputs_count == 2
|
||||
assert count_previewable_outputs(outputs) == 1
|
||||
|
||||
def test_save_text_file_output_is_previewable_by_extension(self):
|
||||
"""SaveText (comfy_extras/nodes_text.py) emits its saved file under a
|
||||
'files' media type via ui.SavedResult: {'files': [{'filename':
|
||||
'..._00001.txt', 'subfolder': ..., 'type': 'output'}]}. The .txt
|
||||
extension makes it previewable even though 'files' itself isn't a
|
||||
previewable media type."""
|
||||
outputs = {
|
||||
'node1': {
|
||||
'files': [{'filename': 'ComfyUI_00001.txt', 'subfolder': '', 'type': 'output'}]
|
||||
}
|
||||
}
|
||||
assert count_previewable_outputs(outputs) == 1
|
||||
|
||||
def test_preview_any_text_tuple_not_counted(self):
|
||||
"""PreviewAny (comfy_extras/nodes_preview_any.py) emits only
|
||||
{'text': (value,)} with no saved file. Since the value is a tuple,
|
||||
not a list, it is excluded from both outputs_count and
|
||||
previewable_outputs_count — matching get_outputs_summary()."""
|
||||
outputs = {
|
||||
'node1': {'text': ('some previewed value',)}
|
||||
}
|
||||
outputs_count, _ = get_outputs_summary(outputs)
|
||||
assert outputs_count == 0
|
||||
assert count_previewable_outputs(outputs) == 0
|
||||
|
||||
def test_string_3d_filename_previewable(self):
|
||||
"""String 3D filenames (e.g. Preview3D) normalize into a previewable
|
||||
item just like they do for outputs_count."""
|
||||
outputs = {
|
||||
'node1': {'result': ['preview3d_abc123.glb', None]}
|
||||
}
|
||||
assert count_previewable_outputs(outputs) == 1
|
||||
|
||||
|
||||
class TestHas3DExtension:
|
||||
"""Unit tests for has_3d_extension()"""
|
||||
|
||||
|
|
@ -447,6 +521,7 @@ class TestNormalizeQueueItem:
|
|||
assert 'execution_error' not in job
|
||||
assert 'preview_output' not in job
|
||||
assert job['outputs_count'] == 0
|
||||
assert job['previewable_outputs_count'] == 0
|
||||
assert job['workflow_id'] == 'workflow-abc'
|
||||
|
||||
|
||||
|
|
@ -635,6 +710,54 @@ class TestNormalizeHistoryItem:
|
|||
{'filename': 'photo.png', 'type': 'output', 'subfolder': ''},
|
||||
]
|
||||
|
||||
def test_previewable_outputs_count_excludes_non_previewable_outputs(self):
|
||||
"""Regression test for the Media Assets badge overcount: a job with an
|
||||
image (SaveImage) and a SaveLatent output should report previewable_
|
||||
outputs_count == 1 while outputs_count == 2, so the frontend badge
|
||||
(once switched to previewable_outputs_count) matches what the
|
||||
expanded asset view actually renders."""
|
||||
history_item = {
|
||||
'prompt': (
|
||||
5,
|
||||
'prompt-mixed',
|
||||
{'nodes': {}},
|
||||
{'create_time': 1234567890},
|
||||
['node1', 'node2'],
|
||||
),
|
||||
'status': {'status_str': 'success', 'completed': True, 'messages': []},
|
||||
'outputs': {
|
||||
'node1': {
|
||||
'images': [{'filename': 'ComfyUI_00001_.png', 'subfolder': '', 'type': 'output'}]
|
||||
},
|
||||
'node2': {
|
||||
'latents': [{'filename': 'ComfyUI_00001_.latent', 'subfolder': '', 'type': 'output'}]
|
||||
},
|
||||
},
|
||||
}
|
||||
job = normalize_history_item('prompt-mixed', history_item)
|
||||
|
||||
assert job['outputs_count'] == 2
|
||||
assert job['previewable_outputs_count'] == 1
|
||||
|
||||
def test_previewable_outputs_count_zero_pruned_by_prune_dict(self):
|
||||
"""A job with no outputs at all should still report both counts as 0,
|
||||
not omit the field (prune_dict only strips None, not 0)."""
|
||||
history_item = {
|
||||
'prompt': (
|
||||
5,
|
||||
'prompt-empty',
|
||||
{'nodes': {}},
|
||||
{'create_time': 1234567890},
|
||||
['node1'],
|
||||
),
|
||||
'status': {'status_str': 'success', 'completed': True, 'messages': []},
|
||||
'outputs': {},
|
||||
}
|
||||
job = normalize_history_item('prompt-empty', history_item)
|
||||
|
||||
assert job['outputs_count'] == 0
|
||||
assert job['previewable_outputs_count'] == 0
|
||||
|
||||
|
||||
class TestNormalizeOutputItem:
|
||||
"""Unit tests for normalize_output_item()"""
|
||||
|
|
|
|||
Loading…
Reference in New Issue