zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-25 18:32:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make a context manager for cast_bias_weight and use it. (#14750)

Browse Source
This commit is contained in:
blepping
2026-08-08 20:38:34 -06:00
committed by GitHub
parent 40e46c7110
commit cbbc9dab1f
4 changed files with 124 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
comfy/controlnet.py
View File

@@ -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__(