From d59690105518ab5eef50690028a7b0fe0d793631 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Wed, 13 Aug 2025 10:00:37 +0300 Subject: [PATCH] Allow stand-in with VACE Don't know how well it can actually work, has some effect --- wanvideo/modules/model.py | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index fb1f72d..88ae5df 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -19,17 +19,18 @@ except: from .attention import attention import numpy as np -__all__ = ['WanModel'] from tqdm import tqdm import gc -from comfy import model_management as mm + from ...utils import log, get_module_memory_mb from ...cache_methods.cache_methods import TeaCacheState, MagCacheState, EasyCacheState, relative_l1_distance from ...multitalk.multitalk import get_attn_map_with_target from ...echoshot.echoshot import rope_apply_z, rope_apply_c, rope_apply_echoshot -from comfy.model_management import get_torch_device, get_autocast_device, get_offload_stream +__all__ = ['WanModel'] + +from comfy import model_management as mm from comfy.ldm.flux.math import apply_rope as apply_rope_comfy def apply_rope_comfy_chunked(xq, xk, freqs_cis, num_chunks=4): @@ -156,7 +157,7 @@ def rope_params(max_seq_len, dim, theta=10000, L_test=25, k=0): freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs -@torch.autocast(device_type=get_autocast_device(get_torch_device()), enabled=False) +@torch.autocast(device_type=mm.get_autocast_device(mm.get_torch_device()), enabled=False) @torch.compiler.disable() def rope_apply(x, grid_sizes, freqs): n, c = x.size(2), x.size(3) // 2 @@ -739,7 +740,7 @@ class WanAttentionBlock(nn.Module): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.get_mod(e.to(x.device)) input_x = self.modulate(self.norm1(x), shift_msa, scale_msa) - if e_ip is not None: + if x_ip is not None: shift_msa_ip, scale_msa_ip, gate_msa_ip, shift_mlp_ip, scale_mlp_ip, gate_mlp_ip = self.get_mod(e_ip.to(x.device)) input_x_ip = self.modulate(self.norm1(x_ip), shift_msa_ip, scale_msa_ip) self.self_attn.cond_size = input_x_ip.shape[1] @@ -862,9 +863,8 @@ class WanAttentionBlock(nn.Module): x_ip = x_ip.addcmul(y_ip, gate_msa_ip) y_ip = self.ffn(torch.addcmul(shift_mlp_ip, self.norm2(x_ip), 1 + scale_mlp_ip)) x_ip = x_ip.addcmul(y_ip, gate_mlp_ip) - return x, x_ip - - return x + + return x, x_ip def cross_attn_ffn(self, x, context, grid_sizes, shift_mlp, scale_mlp, gate_mlp, clip_embed, @@ -986,14 +986,14 @@ class BaseWanAttentionBlock(WanAttentionBlock): self.block_id = block_id def forward(self, x, vace_hints=None, vace_context_scale=[1.0], **kwargs): - x = super().forward(x, **kwargs) + x, x_ip = super().forward(x, **kwargs) if vace_hints is None: - return x + return x, x_ip if self.block_id is not None: for i in range(len(vace_hints)): x.add_(vace_hints[i][self.block_id].to(x.device), alpha=vace_context_scale[i]) - return x + return x, x_ip class Head(nn.Module): @@ -1405,7 +1405,7 @@ class WanModel(torch.nn.Module): else: c_processed = current_c - c_processed = block.forward(c_processed, **kwargs) + c_processed, _ = block.forward(c_processed, **kwargs) # Store skip connection c_skip = block.after_proj(c_processed) @@ -1963,10 +1963,7 @@ class WanModel(torch.nn.Module): if b in self.slg_blocks and is_uncond: if self.slg_start_percent <= current_step_percentage <= self.slg_end_percent: continue - if x_ip is None: - x = block(x,**kwargs) - else: - x, x_ip = block(x, x_ip=x_ip, **kwargs) + x, x_ip = block(x, x_ip=x_ip, **kwargs) if self.block_swap_debug: compute_end = time.perf_counter() compute_time = compute_end - compute_start