From ac8122d089757da249f1676ef6bb0ae0f5ad7c58 Mon Sep 17 00:00:00 2001 From: drozbay <17261091+drozbay@users.noreply.github.com> Date: Sat, 4 Jul 2026 12:14:18 -0600 Subject: [PATCH] Fix attention patch leak and shape detection in Wan Uni3C controlnet --- comfy/ldm/wan/uni3c.py | 15 +++++------- comfy_extras/nodes_model_patch.py | 38 ++++++++++++++++++++++--------- 2 files changed, 33 insertions(+), 20 deletions(-) diff --git a/comfy/ldm/wan/uni3c.py b/comfy/ldm/wan/uni3c.py index 850cb1b39..827ad2339 100644 --- a/comfy/ldm/wan/uni3c.py +++ b/comfy/ldm/wan/uni3c.py @@ -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 diff --git a/comfy_extras/nodes_model_patch.py b/comfy_extras/nodes_model_patch.py index 575cb3377..55c23d1cf 100644 --- a/comfy_extras/nodes_model_patch.py +++ b/comfy_extras/nodes_model_patch.py @@ -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