diff --git a/invokeai/backend/flux/controlnet/zero_module.py b/invokeai/backend/flux/controlnet/zero_module.py index 53a21861a93..b3a0f88b17b 100644 --- a/invokeai/backend/flux/controlnet/zero_module.py +++ b/invokeai/backend/flux/controlnet/zero_module.py @@ -7,6 +7,10 @@ def zero_module(module: T) -> T: """Initialize the parameters of a module to zero.""" - for p in module.parameters(): - torch.nn.init.zeros_(p) + # Vectorized zero-ing of parameters for efficiency + params = list(module.parameters()) + if params: + with torch.no_grad(): + for p in params: + p.zero_() return module