From e724bd542cc7623f96db56d47600fdb7dadd06aa Mon Sep 17 00:00:00 2001 From: CrazyBoyM Date: Thu, 8 Aug 2024 00:03:44 +0800 Subject: [PATCH] add control_weight support for HunyuanDit Controlnet --- comfy/ldm/hydit/controlnet.py | 3 +++ comfy_extras/nodes_hunyuan.py | 4 +++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/comfy/ldm/hydit/controlnet.py b/comfy/ldm/hydit/controlnet.py index c1ae18cf6..6c5f79171 100644 --- a/comfy/ldm/hydit/controlnet.py +++ b/comfy/ldm/hydit/controlnet.py @@ -240,6 +240,7 @@ class HunYuanControlNet(nn.Module): text_embedding_mask_t5=None, image_meta_size=None, style=None, + control_weight=1.0, transformer_options=None, **kwarg, ): @@ -336,4 +337,6 @@ 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] + return {"output": controls} diff --git a/comfy_extras/nodes_hunyuan.py b/comfy_extras/nodes_hunyuan.py index babb2e80c..4f2ccfe90 100644 --- a/comfy_extras/nodes_hunyuan.py +++ b/comfy_extras/nodes_hunyuan.py @@ -29,6 +29,7 @@ class ControlNetApplyAdvancedHunYuan: "image": ("IMAGE", ), "vae": ("VAE", ), "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "control_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.001}), "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}) }} @@ -39,7 +40,7 @@ class ControlNetApplyAdvancedHunYuan: CATEGORY = "conditioning/controlnet" - def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None): + def apply_controlnet(self, positive, negative, control_net, image, strength, control_weight, start_percent, end_percent, vae=None): if strength == 0: return (positive, negative) @@ -57,6 +58,7 @@ class ControlNetApplyAdvancedHunYuan: c_net = cnets[prev_cnet] else: c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae) + c_net.set_extra_arg('control_weight', control_weight) c_net.set_previous_controlnet(prev_cnet) cnets[prev_cnet] = c_net