Add Flux model support for InstantX style controlnet residuals

This commit is contained in:
eesahe 2024-08-17 18:11:08 +00:00
parent 14af129c55
commit d815648d28

View File

@ -119,14 +119,30 @@ class Flux(nn.Module):
if control is not None: #Controlnet if control is not None: #Controlnet
control_o = control.get("output") control_o = control.get("output")
if isinstance(control_o[0], list):
# InstantX format: [double_block_residuals, single_block_residuals]
double_block_residuals, single_block_residuals = control_o
interval_control = int(torch.ceil(torch.tensor(len(self.double_blocks) / len(double_block_residuals))))
index = i // interval_control
img += double_block_residuals[index]
else:
# XLabs format: list of tensors
if i < len(control_o): if i < len(control_o):
add = control_o[i] add = control_o[i]
if add is not None: if add is not None:
img += add img += add
img = torch.cat((txt, img), 1) img = torch.cat((txt, img), 1)
for block in self.single_blocks: for i in range(len(self.single_blocks)):
img = block(img, vec=vec, pe=pe) img = self.single_blocks[i](img, vec=vec, pe=pe)
if control is not None: #Controlnet
control_o = control.get("output")
if isinstance(control_o[0], list): # InstantX format
double_block_residuals, single_block_residuals = control_o
interval_control = int(torch.ceil(torch.tensor(len(self.single_blocks) / len(single_block_residuals))))
index = i // interval_control
img[:, txt.shape[1] :, ...] += single_block_residuals[index]
img = img[:, txt.shape[1] :, ...] img = img[:, txt.shape[1] :, ...]
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)