From 4d55f2e28ced61f27bea94eedb1a5413a656535e Mon Sep 17 00:00:00 2001 From: CrazyBoyM Date: Thu, 8 Aug 2024 00:34:09 +0800 Subject: [PATCH] use control_weights in HunyuanDit Controlnet --- comfy/ldm/hydit/controlnet.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/comfy/ldm/hydit/controlnet.py b/comfy/ldm/hydit/controlnet.py index 6c5f79171..ce9839bfb 100644 --- a/comfy/ldm/hydit/controlnet.py +++ b/comfy/ldm/hydit/controlnet.py @@ -337,6 +337,12 @@ class HunYuanControlNet(nn.Module): x = block(x, c, text_states, freqs_cis_img) controls.append(self.after_proj_list[layer](x)) # zero linear for output - controls = [control * control_weight for control in controls] + control_weights = [1.0 * (control_weight ** float(19 - i)) for i in range(19)] + assert len(control_weight) == len( + controls + ), "control_weights and controls should have the same length" + controls = [ + control * weight for control, weight in zip(controls, control_weights) + ] return {"output": controls}