Fix attention patch leak and shape detection in Wan Uni3C controlnet
This commit is contained in:
drozbay
parent
2c01c141b3
commit
ac8122d089
|
|
@ -45,9 +45,9 @@ class Uni3CAttentionBlock(nn.Module):
|
|||
operations.Linear(dim, ffn_dim, device=device, dtype=dtype), nn.GELU(approximate='tanh'),
|
||||
operations.Linear(ffn_dim, dim, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, x, temb, freqs, transformer_options={}):
|
||||
def forward(self, x, temb, freqs):
|
||||
norm_x, gate_msa = self.norm1(x, temb)
|
||||
x = x + gate_msa * self.self_attn(norm_x, freqs, transformer_options=transformer_options)
|
||||
x = x + gate_msa * self.self_attn(norm_x, freqs)
|
||||
norm_x, gate_ff = self.norm2(x, temb)
|
||||
x = x + gate_ff * self.ffn(norm_x)
|
||||
return x
|
||||
|
|
@ -59,14 +59,10 @@ class MaskCamEmbed(nn.Module):
|
|||
add_channels=7,
|
||||
mid_channels=256,
|
||||
conv_out_dim=5120,
|
||||
interp=False,
|
||||
device=None, dtype=None, operations=None
|
||||
):
|
||||
super().__init__()
|
||||
if interp:
|
||||
self.mask_padding = [0, 0, 0, 0, 3, 3] # first and last frame conditioning
|
||||
else:
|
||||
self.mask_padding = [0, 0, 0, 0, 3, 0] # first frame conditioning
|
||||
self.mask_padding = [0, 0, 0, 0, 3, 0] # first frame conditioning
|
||||
self.mask_proj = nn.Sequential(
|
||||
operations.Conv3d(add_channels, mid_channels, kernel_size=(4, 8, 8), stride=(4, 8, 8), device=device, dtype=dtype),
|
||||
operations.GroupNorm(mid_channels // 8, mid_channels, device=device, dtype=dtype),
|
||||
|
|
@ -129,6 +125,7 @@ class WanUni3CControlnet(nn.Module):
|
|||
return freqs
|
||||
|
||||
def process_input(self, control_input, render_mask=None, camera_embedding=None):
|
||||
# render_mask/camera_embedding are the checkpoint's extra conditioning path, not wired up yet
|
||||
hidden = self.controlnet_patch_embedding(control_input.float()).to(control_input.dtype)
|
||||
t_len, h_len, w_len = hidden.shape[2:]
|
||||
freqs = self.rope_encode(t_len, h_len, w_len, device=hidden.device, dtype=hidden.dtype)
|
||||
|
|
@ -146,7 +143,7 @@ class WanUni3CControlnet(nn.Module):
|
|||
hidden = self.proj_in(hidden)
|
||||
return hidden, freqs
|
||||
|
||||
def forward_block(self, block_index, hidden, temb, freqs, transformer_options={}):
|
||||
hidden = self.controlnet_blocks[block_index](hidden, temb, freqs, transformer_options=transformer_options)
|
||||
def forward_block(self, block_index, hidden, temb, freqs):
|
||||
hidden = self.controlnet_blocks[block_index](hidden, temb, freqs)
|
||||
residual = self.proj_out[block_index](hidden)
|
||||
return hidden, residual
|
||||
|
|
|
|||
|
|
@ -275,13 +275,18 @@ class ModelPatchLoader:
|
|||
sd = converted_sd
|
||||
|
||||
num_layers = sum(1 for k in sd if k.startswith("proj_out.") and k.endswith(".weight"))
|
||||
conv_out_dim = sd["controlnet_patch_embedding.weight"].shape[0]
|
||||
if "proj_in.weight" in sd:
|
||||
dim = sd["proj_in.weight"].shape[0]
|
||||
else:
|
||||
dim = conv_out_dim
|
||||
model = comfy.ldm.wan.uni3c.WanUni3CControlnet(
|
||||
in_channels=sd["controlnet_patch_embedding.weight"].shape[1],
|
||||
conv_out_dim=sd["controlnet_patch_embedding.weight"].shape[0],
|
||||
dim=sd["proj_out.0.weight"].shape[1],
|
||||
conv_out_dim=conv_out_dim,
|
||||
dim=dim,
|
||||
ffn_dim=sd["controlnet_blocks.0.ffn.0.bias"].shape[0],
|
||||
num_layers=num_layers,
|
||||
time_embed_dim=sd["controlnet_blocks.0.norm1.linear.weight"].shape[1],
|
||||
time_embed_dim=conv_out_dim,
|
||||
out_proj_dim=sd["proj_out.0.weight"].shape[0],
|
||||
add_channels=sd["controlnet_mask_embedding.mask_proj.0.weight"].shape[1],
|
||||
mid_channels=sd["controlnet_mask_embedding.mask_proj.0.weight"].shape[0],
|
||||
|
|
@ -548,6 +553,7 @@ class WanUni3CCnetPatch:
|
|||
self.strength = strength
|
||||
self.sigma_start = sigma_start
|
||||
self.sigma_end = sigma_end
|
||||
self.prepared_render = None
|
||||
self.temp_data = None
|
||||
|
||||
def build_controlnet_input(self, x, dtype):
|
||||
|
|
@ -558,9 +564,13 @@ class WanUni3CCnetPatch:
|
|||
pad_shape[1] = 20 - hidden.shape[1]
|
||||
hidden = torch.cat([hidden, torch.zeros(pad_shape, dtype=hidden.dtype, device=hidden.device)], dim=1)
|
||||
|
||||
render = self.render_latent.to(device=hidden.device, dtype=hidden.dtype)
|
||||
if render.shape[2:] != hidden.shape[2:]:
|
||||
render = torch.nn.functional.interpolate(render, size=hidden.shape[2:], mode="trilinear", align_corners=False)
|
||||
render = self.prepared_render
|
||||
if render is None or render.shape[2:] != hidden.shape[2:] or render.device != hidden.device or render.dtype != dtype:
|
||||
render = self.render_latent.to(device=hidden.device)
|
||||
if render.shape[2:] != hidden.shape[2:]:
|
||||
render = torch.nn.functional.interpolate(render, size=hidden.shape[2:], mode="trilinear", align_corners=False)
|
||||
render = render.to(dtype)
|
||||
self.prepared_render = render
|
||||
return torch.cat([hidden, render], dim=1)
|
||||
|
||||
def __call__(self, kwargs):
|
||||
|
|
@ -577,19 +587,24 @@ class WanUni3CCnetPatch:
|
|||
if sigma > self.sigma_start or sigma < self.sigma_end:
|
||||
active = False
|
||||
if active:
|
||||
controlnet_input = self.build_controlnet_input(kwargs.get("x"), img.dtype)
|
||||
temb = kwargs.get("vec")[:1]
|
||||
if temb.ndim == 3:
|
||||
temb = temb[:, 0]
|
||||
hidden, freqs = self.model_patch.model.process_input(controlnet_input)
|
||||
model = self.model_patch.model
|
||||
cnet_dim = model.controlnet_blocks[0].norm1.linear.in_features
|
||||
if temb.shape[-1] != cnet_dim:
|
||||
raise RuntimeError("This Uni3C controlnet expects a Wan model with dim {}, the loaded model has dim {}".format(cnet_dim, temb.shape[-1]))
|
||||
controlnet_input = self.build_controlnet_input(kwargs.get("x"), img.dtype)
|
||||
hidden, freqs = model.process_input(controlnet_input)
|
||||
self.temp_data = (hidden, temb.to(img.dtype), freqs)
|
||||
|
||||
if self.temp_data is not None and block_index < self.model_patch.model.num_layers:
|
||||
num_layers = self.model_patch.model.num_layers
|
||||
if self.temp_data is not None and block_index < num_layers:
|
||||
hidden, temb, freqs = self.temp_data
|
||||
hidden, residual = self.model_patch.model.forward_block(block_index, hidden, temb, freqs, transformer_options=transformer_options)
|
||||
hidden, residual = self.model_patch.model.forward_block(block_index, hidden, temb, freqs)
|
||||
img_offset = kwargs.get("img_offset", 0)
|
||||
img[:, img_offset:img_offset + residual.shape[1]] += residual.to(img.dtype) * self.strength
|
||||
if block_index >= self.model_patch.model.num_layers - 1:
|
||||
if block_index >= num_layers - 1:
|
||||
self.temp_data = None
|
||||
else:
|
||||
self.temp_data = (hidden, temb, freqs)
|
||||
|
|
@ -599,6 +614,7 @@ class WanUni3CCnetPatch:
|
|||
def to(self, device_or_dtype):
|
||||
if isinstance(device_or_dtype, torch.device):
|
||||
self.render_latent = self.render_latent.to(device_or_dtype)
|
||||
self.prepared_render = None
|
||||
self.temp_data = None
|
||||
return self
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue