Browse Source

Add a way to control controlnet strength.

pull/10/head^2
comfyanonymous 2 years ago
parent
commit
6135a21ee8
  1. 7
      comfy/sd.py
  2. 10
      nodes.py

7
comfy/sd.py

@ -331,6 +331,7 @@ class ControlNet:
self.control_model = control_model
self.cond_hint_original = None
self.cond_hint = None
self.strength = 1.0
def get_control(self, x_noisy, t, cond_txt):
if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
@ -340,10 +341,13 @@ class ControlNet:
self.cond_hint = utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(x_noisy.device)
print("set cond_hint", self.cond_hint.shape)
control = self.control_model(x=x_noisy, hint=self.cond_hint, timesteps=t, context=cond_txt)
for x in control:
x *= self.strength
return control
def set_cond_hint(self, cond_hint):
def set_cond_hint(self, cond_hint, strength=1.0):
self.cond_hint_original = cond_hint
self.strength = strength
return self
def cleanup(self):
@ -354,6 +358,7 @@ class ControlNet:
def copy(self):
c = ControlNet(self.control_model)
c.cond_hint_original = self.cond_hint_original
c.strength = self.strength
return c
def load_controlnet(ckpt_path):

10
nodes.py

@ -234,19 +234,23 @@ class ControlNetLoader:
class ControlNetApply:
@classmethod
def INPUT_TYPES(s):
return {"required": {"conditioning": ("CONDITIONING", ), "control_net": ("CONTROL_NET", ), "image": ("IMAGE", )}}
return {"required": {"conditioning": ("CONDITIONING", ),
"control_net": ("CONTROL_NET", ),
"image": ("IMAGE", ),
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01})
}}
RETURN_TYPES = ("CONDITIONING",)
FUNCTION = "apply_controlnet"
CATEGORY = "conditioning"
def apply_controlnet(self, conditioning, control_net, image):
def apply_controlnet(self, conditioning, control_net, image, strength):
c = []
control_hint = image.movedim(-1,1)
print(control_hint.shape)
for t in conditioning:
n = [t[0], t[1].copy()]
n[1]['control'] = control_net.copy().set_cond_hint(control_hint)
n[1]['control'] = control_net.copy().set_cond_hint(control_hint, strength)
c.append(n)
return (c, )

Loading…
Cancel
Save