nodes/kijai-ComfyUI-WanVideoWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Some cleanup and allow full block swap

Browse Source
This commit is contained in:
kijai
2025-10-26 23:05:17 +02:00
parent 8ad7e50f33
commit a0bdf20817
2 changed files with 13 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes_model_loading.py
+1 -1
View File
@@ -278,7 +278,7 @@ class WanVideoBlockSwap:
def INPUT_TYPES(s):
return {
"required": {
"blocks_to_swap": ("INT", {"default": 20, "min": 0, "max": 40, "step": 1, "tooltip": "Number of transformer blocks to swap, the 14B model has 40, while the 1.3B model has 30 blocks"}),
"blocks_to_swap": ("INT", {"default": 20, "min": 0, "max": 48, "step": 1, "tooltip": "Number of transformer blocks to swap, the 14B model has 40, while the 1.3B and 5B models have 30 blocks. LongCat-video has 48 blocks"}),
"offload_img_emb": ("BOOLEAN", {"default": False, "tooltip": "Offload img_emb to offload_device"}),
"offload_txt_emb": ("BOOLEAN", {"default": False, "tooltip": "Offload time_emb to offload_device"}),
},
wanvideo/modules/model.py
+12 -13
View File
@@ -1232,19 +1232,18 @@ class WanAttentionBlock(nn.Module):
full_k = torch.cat([k, k_ip], dim=1)
full_v = torch.cat([v, v_ip], dim=1)
y = self.self_attn.forward(q, full_k, full_v, seq_lens)
elif is_longcat:
if num_cond_latents is not None and num_cond_latents > 0:
num_cond_latents_thw = num_cond_latents * (N // grid_sizes[0][0])
# process the condition tokens
q_cond = q[:, :num_cond_latents_thw].contiguous()
k_cond = k[:, :num_cond_latents_thw].contiguous()
v_cond = v[:, :num_cond_latents_thw].contiguous()
x_cond = self.self_attn.forward(q_cond, k_cond, v_cond, seq_lens)
# process the noise tokens
q_noise = q[:, num_cond_latents_thw:].contiguous()
x_noise = self.self_attn.forward(q_noise, k, v, seq_lens)
# merge x_cond and x_noise
y = torch.cat([x_cond, x_noise], dim=1).contiguous()
elif is_longcat and num_cond_latents is not None and num_cond_latents > 0:
num_cond_latents_thw = num_cond_latents * (N // grid_sizes[0][0])
# process the condition tokens
x_cond = self.self_attn.forward(
q[:, :num_cond_latents_thw].contiguous(),
k[:, :num_cond_latents_thw].contiguous(),
v[:, :num_cond_latents_thw].contiguous(),
seq_lens)
# process the noise tokens
x_noise = self.self_attn.forward(q[:, num_cond_latents_thw:].contiguous(), k, v, seq_lens)
# merge x_cond and x_noise
y = torch.cat([x_cond, x_noise], dim=1).contiguous()
else:
y = self.self_attn.forward(q, k, v, seq_lens, lynx_ref_feature=lynx_ref_feature, lynx_ref_scale=lynx_ref_scale)