andreas/ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add ControlNet support.

Browse Source
This commit is contained in:
comfyanonymous
2023-02-16 10:38:08 -05:00
parent bc69fb5245
commit 4efa67fa12
9 changed files with 580 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
comfy/ldm/models/diffusion/ddpm.py
View File

@@ -1320,12 +1320,12 @@ class DiffusionWrapper(torch.nn.Module):
self.conditioning_key = conditioning_key
assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm', 'hybrid-adm', 'crossattn-adm']
def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None):
def forward(self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None, control=None):
if self.conditioning_key is None:
out = self.diffusion_model(x, t)
out = self.diffusion_model(x, t, control=control)
elif self.conditioning_key == 'concat':
xc = torch.cat([x] + c_concat, dim=1)
out = self.diffusion_model(xc, t)
out = self.diffusion_model(xc, t, control=control)
elif self.conditioning_key == 'crossattn':
if not self.sequential_cross_attn:
cc = torch.cat(c_crossattn, 1)
@@ -1335,25 +1335,25 @@ class DiffusionWrapper(torch.nn.Module):
# TorchScript changes names of the arguments
# with argument cc defined as context=cc scripted model will produce
# an error: RuntimeError: forward() is missing value for argument 'argument_3'.
out = self.scripted_diffusion_model(x, t, cc)
out = self.scripted_diffusion_model(x, t, cc, control=control)
else:
out = self.diffusion_model(x, t, context=cc)
out = self.diffusion_model(x, t, context=cc, control=control)
elif self.conditioning_key == 'hybrid':
xc = torch.cat([x] + c_concat, dim=1)
cc = torch.cat(c_crossattn, 1)
out = self.diffusion_model(xc, t, context=cc)
out = self.diffusion_model(xc, t, context=cc, control=control)
elif self.conditioning_key == 'hybrid-adm':
assert c_adm is not None
xc = torch.cat([x] + c_concat, dim=1)
cc = torch.cat(c_crossattn, 1)
out = self.diffusion_model(xc, t, context=cc, y=c_adm)
out = self.diffusion_model(xc, t, context=cc, y=c_adm, control=control)
elif self.conditioning_key == 'crossattn-adm':
assert c_adm is not None
cc = torch.cat(c_crossattn, 1)
out = self.diffusion_model(x, t, context=cc, y=c_adm)
out = self.diffusion_model(x, t, context=cc, y=c_adm, control=control)
elif self.conditioning_key == 'adm':
cc = c_crossattn[0]
out = self.diffusion_model(x, t, y=cc)
out = self.diffusion_model(x, t, y=cc, control=control)
else:
raise NotImplementedError()

10
comfy/ldm/modules/diffusionmodules/openaimodel.py
View File

@@ -753,7 +753,7 @@ class UNetModel(nn.Module):
self.middle_block.apply(convert_module_to_f32)
self.output_blocks.apply(convert_module_to_f32)
def forward(self, x, timesteps=None, context=None, y=None,**kwargs):
def forward(self, x, timesteps=None, context=None, y=None, control=None, **kwargs):
"""
Apply the model to an input batch.
:param x: an [N x C x ...] Tensor of inputs.
@@ -778,8 +778,14 @@ class UNetModel(nn.Module):
h = module(h, emb, context)
hs.append(h)
h = self.middle_block(h, emb, context)
if control is not None:
h += control.pop()
for module in self.output_blocks:
h = th.cat([h, hs.pop()], dim=1)
hsp = hs.pop()
if control is not None:
hsp += control.pop()
h = th.cat([h, hsp], dim=1)
h = module(h, emb, context)
h = h.type(x.dtype)
if self.predict_codebook_ids: